from PIL import Image
import struct
import numpy as np
import matplotlib.pyplot as plt

class DepthColorProcessor:
    """
    A single class that handles reading depth and color frames from binary files,
    normalizing depth data to 8-bit (for visualization), and displaying or
    continuously streaming the results.

    Usage Example:
        dcp = DepthColorProcessor()
        dcp.main(depth_file="depth_data.bin", color_file="color_data.bin")
    """

    def __init__(self, interactive_plots=True):
        """
        :param interactive_plots: If True, enables plt.ion() to render plots non-blocking.
        """
        self.interactive_plots = interactive_plots
        if self.interactive_plots:
            plt.ion()

    def read_depth_color(self, depth_bin, color_bin):
        """
        Reads a single-frame depth_bin and color_bin.

        The returned depth_array is 16-bit per pixel (shape: [height, width]),
        which is *not* a typical image format. The returned color_array is
        8-bit per channel (shape: [height, width, 3]) in RGB format.
        """

        # --- Read Depth File ---
        with open(depth_bin, "rb") as f:
            header_data = f.read(24)
            magic, w, h, depth_format, bpp, frame_count = struct.unpack("<6I", header_data)

            if magic != 0x44455054:  # 'DEPT'
                raise ValueError(f"Unexpected magic number in {depth_bin}: {hex(magic)}")
            if bpp != 2:
                raise ValueError(f"Expected 2 bytes_per_pixel for depth, got {bpp}")

            depth_frame_bytes = w * h * bpp
            depth_buffer = f.read(depth_frame_bytes)
            if len(depth_buffer) != depth_frame_bytes:
                raise ValueError(f"Expected {depth_frame_bytes} bytes, "
                                 f"found {len(depth_buffer)} in {depth_bin}")
            depth_data = np.frombuffer(depth_buffer, dtype=np.uint16).reshape((h, w))
        # --- Read Color File ---
        with open(color_bin, "rb") as f:
            header_data = f.read(24)
            magic, w2, h2, color_format, bpp, frame_count = struct.unpack("<6I", header_data)

            if magic != 0x434F4C52:  # 'COLR'
                raise ValueError(f"Unexpected magic number in {color_bin}: {hex(magic)}")
            if bpp != 3:
                raise ValueError(f"Expected 3 bytes_per_pixel for color, got {bpp}")

            if (w, h) != (w2, h2):
                raise ValueError(
                    "Depth and color dimensions differ. "
                    f"Depth: {w}x{h}, Color: {w2}x{h2}"
                )

            color_frame_bytes = w2 * h2 * bpp
            color_buffer = f.read(color_frame_bytes)
            if len(color_buffer) != color_frame_bytes:
                raise ValueError(f"Expected {color_frame_bytes} bytes, "
                                 f"found {len(color_buffer)} in {color_bin}")

            color_data = np.frombuffer(color_buffer, dtype=np.uint8).reshape((h2, w2, 3))

        return depth_data, color_data

    def normalize_to_8bit(self, depth_frame):
        """
        Normalize 16-bit depth data into an 8-bit range [0..255].
        This is only for quick visualization; depth is not a typical 8-bit image.
        """
        depth_min, depth_max = depth_frame.min(), depth_frame.max()
        if depth_min == depth_max:
            return np.zeros_like(depth_frame, dtype=np.uint8)
        normalized = (depth_frame - depth_min) / float(depth_max - depth_min)
        return (normalized * 255).astype(np.uint8)

    def display_depth_frame(self, depth_frame_8u):
        """
        Display a *single* 8-bit visualization of the depth data.
        """
        plt.clf()  # Clear the previous image
        plt.imshow(depth_frame_8u, cmap='gray')
        plt.title("Depth Frame (visualized as 8-bit grayscale)")
        plt.axis('off')
        plt.pause(0.001)

    def display_color_frame(self, color_frame):
        """
        Display the color frame (RGB) using matplotlib.
        We use a separate figure so it doesn't overwrite the depth display.
        """
        plt.figure()
        plt.imshow(color_frame)
        plt.title("Color Frame (RGB)")
        plt.axis('off')
        plt.pause(0.001)

    def display_depth_color(self, depth_file="depth_data.bin", color_file="color_data.bin"):
        """
        Main showcase method:
          Reads depth & color data from the specified files, prints out info,
          and displays both frames (depth in grayscale, color in RGB).
        """
        depth_data, color_data = self.read_depth_color(depth_file, color_file)
        print("Depth array shape:", depth_data.shape, "dtype:", depth_data.dtype)
        print("Color array shape:", color_data.shape, "dtype:", color_data.dtype)

        depth_8u = self.normalize_to_8bit(depth_data)
        self.display_depth_frame(depth_8u)
        self.display_color_frame(color_data)

        # Keep plots open for a moment, especially if interactive mode is not used:
        if not self.interactive_plots:
            plt.show()

    def stream_color_continuous(self, color_bin="color_data.bin", output_png_path="color.png"):
        """
        Continuously reads from color_bin, generating an RGB image each time.
        Saves the image to output_png_path. Press Ctrl+C to exit.
        """
        while True:
            with open(color_bin, "rb") as f:
                header_data = f.read(24)
                magic, w, h, color_format, bpp, frame_count = struct.unpack("<6I", header_data)
                if magic != 0x434F4C52:  # 'COLR'
                    raise ValueError(f"Unexpected magic number in {color_bin}: {hex(magic)}")
                if bpp != 3:
                    raise ValueError(f"Expected 3 bytes_per_pixel for color, got {bpp}")

                color_frame_bytes = w * h * bpp
                color_buffer = f.read(color_frame_bytes)
                if len(color_buffer) != color_frame_bytes:
                    raise ValueError(f"Expected {color_frame_bytes} bytes, "
                                     f"found {len(color_buffer)} in {color_bin}")

                color_data = np.frombuffer(color_buffer, dtype=np.uint8).reshape((h, w, 3))
            # Create an RGB image and write to PNG
            img = Image.fromarray(color_data, mode="RGB")
            img.save(output_png_path)
            # Optional: add a small sleep if desired to reduce CPU usage
            # time.sleep(0.1)

    def depth_bin_to_png(self, depth_bin="depth_data.bin", out_png="depth.png"):
        """
        Read a depth .bin file and save it as a grayscale PNG.
        """
        with open(depth_bin, "rb") as f:
            header_data = f.read(24)
            magic, w, h, depth_format, bpp, frame_count = struct.unpack("<6I", header_data)

            if magic != 0x44455054:  # 'DEPT'
                raise ValueError(f"Unexpected magic number in {depth_bin}: {hex(magic)}")
            if bpp != 2:
                raise ValueError(f"Expected 2 bytes_per_pixel for depth, got {bpp}")

            depth_frame_bytes = w * h * bpp
            depth_buffer = f.read(depth_frame_bytes)
            if len(depth_buffer) != depth_frame_bytes:
                raise ValueError(
                    f"Expected {depth_frame_bytes} bytes, "
                    f"found {len(depth_buffer)} in {depth_bin}"
                )

            depth_data = np.frombuffer(depth_buffer, dtype=np.uint16).reshape((h, w))

        # Normalize to 8-bit
        depth_min = depth_data.min()
        depth_max = depth_data.max()
        if depth_min == depth_max:
            depth_8u = np.zeros_like(depth_data, dtype=np.uint8)
        else:
            normalized = (depth_data - depth_min) / float(depth_max - depth_min)
            depth_8u = (normalized * 255).astype(np.uint8)

        # Convert to a PIL image and save as PNG
        img = Image.fromarray(depth_8u, mode="L")
        img.save(out_png)
        return depth_8u

    def color_bin_to_png(self, color_bin="color_data.bin", out_png="color.png"):
        """
        Read a color .bin file and save it as an RGB PNG.
        """
        with open(color_bin, "rb") as f:
            header_data = f.read(24)
            magic, w, h, color_format, bpp, frame_count = struct.unpack("<6I", header_data)

            if magic != 0x434F4C52:  # 'COLR'
                raise ValueError(f"Unexpected magic number in {color_bin}: {hex(magic)}")
            if bpp != 3:
                raise ValueError(f"Expected 3 bytes_per_pixel for color, got {bpp}")

            color_frame_bytes = w * h * bpp
            color_buffer = f.read(color_frame_bytes)
            if len(color_buffer) != color_frame_bytes:
                raise ValueError(
                    f"Expected {color_frame_bytes} bytes, "
                    f"found {len(color_buffer)} in {color_bin}"
                )

            color_data = np.frombuffer(color_buffer, dtype=np.uint8).reshape((h, w, 3))

        # Convert to a PIL image and save as PNG
        img = Image.fromarray(color_data, mode="RGB")
        img.save(out_png)
        return color_data

    def getDepthByXY(self, depth_file_path, x, y):
        """
        Reads the header and single-frame data from depth_file_path,
        then returns the depth value at the (x, y) coordinate.
        NOTE: (x, y) must lie within the valid image dimensions.
        """
        with open(depth_file_path, "rb") as f:
            # Read the file header (24 bytes)
            header_data = f.read(24)
            magic, w, h, depth_format, bpp, frame_count = struct.unpack("<6I", header_data)
            # Validate the header
            if magic != 0x44455054:  # 'DEPT'
                raise ValueError(f"Unexpected magic number in {depth_file_path}: {hex(magic)}")
            if bpp != 2:
                raise ValueError(f"Expected 2 bytes_per_pixel for depth, got {bpp}")
            # Check that (x, y) is in range
            if x < 0 or x >= w or y < 0 or y >= h:
                raise ValueError(f"Coordinates (x={x}, y={y}) out of range. Valid range is x in [0,{w-1}], y in [0,{h-1}].")
            # Read the raw depth frame
            depth_frame_bytes = w * h * bpp
            depth_buffer = f.read(depth_frame_bytes)
            if len(depth_buffer) != depth_frame_bytes:
                raise ValueError(
                    f"Expected {depth_frame_bytes} bytes, found {len(depth_buffer)} in {depth_file_path}"
                )
            # Convert to numpy array and reshape to (height, width)
            depth_data = np.frombuffer(depth_buffer, dtype=np.uint16).reshape((h, w))
        return depth_data[y, x]

def test_reconstruct_depth(
    original_depth_file="depth_data.bin",
    reconstructed_depth_file="depth_data_re.bin",
    reconstructed_png_file="output_depth_re.png"
):
    """
    Test that queries every pixel from original_depth_file
    via getDepthByXY, reconstructs a new depth BIN file,
    and then converts that new file to a PNG.
    """

    dcp = DepthColorProcessor(interactive_plots=False)

    # 1. Read the header of the original depth file to get dimensions, etc.
    with open(original_depth_file, "rb") as f:
        header_data = f.read(24)
        magic, w, h, depth_format, bpp, frame_count = struct.unpack("<6I", header_data)
        # Make sure it's a valid depth BIN
        if magic != 0x44455054:  # 'DEPT'
            raise ValueError(f"Unexpected magic number in {original_depth_file}: {hex(magic)}")
        if bpp != 2:
            raise ValueError(f"Expected 2 bytes_per_pixel for depth, got {bpp}")

    # 2. Allocate a new numpy array to hold the reconstructed depth data
    reconstructed_data = np.zeros((h, w), dtype=np.uint16)

    # 3. For each pixel, read its value via getDepthByXY, store in reconstructed_data
    for y in range(h):
        for x in range(w):
            depth_val = dcp.getDepthByXY(original_depth_file, x, y)
            reconstructed_data[y, x] = depth_val
        print(f"Processed row {y} of {h}")

    # 4. Write the new header + reconstructed buffer to a new .bin file
    #    We'll replicate the original header, but the pixel bytes come
    #    from our new array.
    with open(reconstructed_depth_file, "wb") as f_out:
        # Write the same 24-byte header
        header_packed = struct.pack("<6I", magic, w, h, depth_format, bpp, frame_count)
        f_out.write(header_packed)
        # Write the actual depth data
        f_out.write(reconstructed_data.tobytes())

    # 5. Use depth_bin_to_png to convert the new file to a PNG
    dcp.depth_bin_to_png(reconstructed_depth_file, reconstructed_png_file)


if __name__ == "__main__":
    # test_reconstruct_depth()
    dcp = DepthColorProcessor(interactive_plots=False)
    dcp.display_depth_color("depth_data.bin", "color_data.bin") 
    dcp.depth_bin_to_png("depth_data.bin", "output_depth.png")
    dcp.color_bin_to_png("color_data.bin", "output_color.png")

